// Copyright (c) 2018 Baidu.com, Inc. All Rights Reserved
// 
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
// 
//     http://www.apache.org/licenses/LICENSE-2.0
// 
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include <gflags/gflags.h>
#include <bthread/bthread.h>
#include <brpc/channel.h>
#include <brpc/controller.h>
#include <braft/raft.h>
#include <braft/util.h>
#include <braft/route_table.h>
#include "block.pb.h"

DEFINE_bool(log_each_request, false, "Print log for each request");
DEFINE_bool(use_bthread, false, "Use bthread to send requests");
DEFINE_int32(block_size, 64 * 1024 * 1024u, "Size of block");
DEFINE_uint64(request_size, 1024, "Size of each requst");
DEFINE_int32(thread_num, 1, "Number of threads sending requests");
DEFINE_int32(timeout_ms, 5000, "Timeout for each request");
DEFINE_int32(write_percentage, 100, "Percentage of fetch_add");
DEFINE_string(conf, "", "Configuration of the raft group");
DEFINE_string(group, "Block", "Id of the replication group");
DEFINE_int32(group_count, 3, "default of group count");
DEFINE_uint64(file_size, 10, "file size #GB");
DEFINE_uint64(chunk_size, 4, "file size #MB");
DEFINE_string(request_mode, "rand", "reuest_mode:rand or seq");

bvar::LatencyRecorder g_latency_recorder("block_client");

class Sender{
  explicit Sender(){}

public:
  void select_leader(int group){
    std::string group_name = "group_" + std::to_string(group);
    while(braft::rtb::select_leader(group_name, _leader[group])!=0){
       butil::Status st = braft::rtb::refresh_leader(
          group_name, FLAGS_timeout_ms);
       if( !st.ok()){
          LOG(WARNING) << "Fail to refresh_leader: " << st;
          bthread_usleep(FLAGS_timeout_ms * 1000L);
       }
    }
    LOG(INFO) << "group" << group << " Leader is: " << *_leader[group];
  }

  void init(){
    _group_count = FLAGS_group_count;
    _channel = new brpc::Channel*[_group_count];
    _leader = new braft::PeerId* [_group_count];
    _stub = new example::BlockService_Stub* [_group_count];
    for (int group = 0; group < _group_count; group++){
        _channel[group]=  new brpc::Channel();
        _leader[group] =  new braft::PeerId();
        select_leader(group);
        if (_channel[group]->Init(_leader[group]->addr, NULL)!=0){
          LOG(ERROR) << "Failed to init channel to " << _leader[group];
          return;
        }
        _stub[group] = new example::BlockService_Stub(_channel[group]);
    }
    _buf = (char*) malloc(FLAGS_request_size);
  }
  
  void send_one(int group, int chunkid, int offset, int size){
    brpc::Controller cntl;
    //example::BlockService_Stub stub(&channel);
    cntl.set_timeout_ms(FLAGS_timeout_ms);
      // Randomly select which request we want send;
    example::BlockRequest request;
    example::BlockResponse response;
    const char* op = NULL;
    if (butil::fast_rand_less_than(100) < (size_t)FLAGS_write_percentage) {
      op = "write";
      request.set_blockid(group);
      request.set_chunkid(chunkid);
      request.set_offset(offset);
      //sprintf(buf, "%d", data++);
      memset(_buf,'a', size);
      cntl.request_attachment().append(_buf, size);
      _stub[group]->write(&cntl, &request, &response, NULL);
    } else {
	    op = "read";
      request.set_blockid(group);
	    request.set_chunkid(chunkid);
      request.set_offset(offset);
	    request.set_size(size);
	   _stub[group]->read(&cntl, &request, &response, NULL);
    }

    if (cntl.Failed()) {
      LOG(WARNING) << "Fail to send request to " << *_leader[group]
        << " : " << cntl.ErrorText();
      // Clear leadership since this RPC failed.
      std::string group_name = "group_" + std::to_string(group);
      braft::rtb::update_leader(group_name, braft::PeerId());
      bthread_usleep(FLAGS_timeout_ms * 1000L);
      select_leader(group);
     }
     if (!response.success()) {
	     LOG(WARNING) << "Fail to send request to " << *_leader[group]
		     << ", redirecting to "
		     << (response.has_redirect()
				     ? response.redirect() : "nowhere");
         std::string group_name = "group_" + std::to_string(group);
	       // Update route table since we have redirect information
	       braft::rtb::update_leader(group_name, response.redirect());
         select_leader(group);
     }
     g_latency_recorder << cntl.latency_us();

     if (FLAGS_log_each_request) {
	     LOG(INFO) << "Received response from " << *_leader[group] 
                     <<" blockid =" << request.blockid()
		     << " chunkid =" << request.chunkid()
		     << " offset =" << request.offset()
		     << " size =" << request.size()
		     << " op=" << op
		     << " request_attachment="
		     << cntl.request_attachment().size()
		     << " data:" << cntl.request_attachment()
		     << " response_attachment="
		     << cntl.response_attachment().size()
		     << " data:"<< cntl.response_attachment()
		     << " latency=" << cntl.latency_us();
	     bthread_usleep(1000L * 1000L);
     }
  }

  int get_group_count(){
	  return _group_count;
  }

  static void* send_test(void* arg){
    uint64_t offset = 0;
    uint64_t file_size = FLAGS_file_size * 1024 * 1024 * 1024;
    uint64_t chunk_size = FLAGS_chunk_size * 1024 * 1024;
    Sender sender;
    sender.init();
    uint64_t group, chunkid, off_in_chunk;
    uint64_t group_count = sender.get_group_count();
          printf("request_mode:%s\n", FLAGS_request_mode.c_str());
    while (!brpc::IsAskedToQuit()){
      if(FLAGS_request_mode =="rand"){ //rand
        offset = butil::fast_rand_less_than(file_size);
      }else{  //sequence
        offset+= FLAGS_request_size;
      }
      chunkid = offset / chunk_size;
      group = chunkid % group_count ;
      off_in_chunk = offset % chunk_size;
            //printf("offset:%lu,chunk_id:%lu,group:%lu, off_in_chunk:%lu,size:%lu\n",offset,chunkid, group, off_in_chunk, FLAGS_request_size);
      sender.send_one(group, chunkid, off_in_chunk, FLAGS_request_size);
    }
    return 0;
  }

private:
  int _group_count;
  brpc::Channel** _channel;
  braft::PeerId** _leader;
  example::BlockService_Stub** _stub;
  int _write_index = 0;
  int _read_index =0;
  char* _buf;
};



int init_conf(){
   for(int i=0; i< FLAGS_group_count; i++){
      std::string group_name = "group_" + std::to_string(i);
      if (braft::rtb::update_configuration(group_name, FLAGS_conf) != 0) {
        LOG(ERROR) << "Fail to register configuration " << FLAGS_conf
                   << " of group " << group_name;
        return -1;
      }

   }
   return 9;
}


int main(int argc, char* argv[]) {
    GFLAGS_NS::ParseCommandLineFlags(&argc, &argv, true);
    butil::AtExitManager exit_manager;

    init_conf();

    std::vector<bthread_t> tids;
    tids.resize(FLAGS_thread_num);
    if (!FLAGS_use_bthread) {
        for (int i = 0; i < FLAGS_thread_num; ++i) {
            if (pthread_create(&tids[i], NULL, Sender::send_test, NULL) != 0) {
                LOG(ERROR) << "Fail to create pthread";
                return -1;
            }
        }
    } else {
        for (int i = 0; i < FLAGS_thread_num; ++i) {
            if (bthread_start_background(&tids[i], NULL, Sender::send_test, NULL) != 0) {
                LOG(ERROR) << "Fail to create bthread";
                return -1;
            }
        }
    }

    while (!brpc::IsAskedToQuit()) {
        sleep(1);
        LOG_IF(INFO, !FLAGS_log_each_request)
                << "Sending Request to "
                << " (" << FLAGS_conf << ')'
                << " at qps=" << g_latency_recorder.qps(1)
                << " at bw(KB)=" << g_latency_recorder.qps(1) * FLAGS_request_size / 1024 
                << " latency=" << g_latency_recorder.latency(1);
    }

    LOG(INFO) << "Block client is going to quit";
    for (int i = 0; i < FLAGS_thread_num; ++i) {
        if (!FLAGS_use_bthread) {
            pthread_join(tids[i], NULL);
        } else {
            bthread_join(tids[i], NULL);
        }
    }

    return 0;
}
